from common.config import get_config_path
from engine.strategy import Strategy
from engine.data.dataloader import D
from engine.backtest import Backtest
import importlib
import os
import pandas as pd
from datetime import datetime

class TaskRunner:
    def __init__(self):
        self._load_tasks()

    def _load_tasks(self):
        import toml
        self.tasks = {}
        task_path = get_config_path() + '/tasks'
        for file in os.listdir(task_path):
            task = toml.load(task_path + '/' + file)
            self.tasks[task['uid']] = task

    def run_by_uid(self, uid):
        if uid in self.tasks.keys():
            config = self.tasks[uid]
            df = self.run(config)
            self.save_result(uid, df)
            return df
        else:
            print('任务配置：{}不存在'.format(uid))

    def save_result(self, uid, df):
        path = get_config_path() + '/results/{}.csv'.format(uid)
        df.to_csv(path)

    def run(self, config: dict):

        df = self.prepare_data(config)
        #df.to_csv(get_config_path() + '/cache.csv')
        s = self.prepare_strategy(config)
        print(s)
        b = Backtest(df=df)
        df = b.run(s)


        return df
        '''
        df['equity'] = (1+df['portfolio'].pct_change()).cumprod()
        print(df)
        df['equity'].plot()
        import matplotlib.pyplot as plt
        plt.show()
        '''
        #df.to_csv(get_config_path() + '/results/{}.csv'.format(s.name))


    def prepare_data(self,config:dict):

        fields = []
        if 'rules' in config.keys():
            rules = config['rules']
            buy_rules = rules['buy_rules']
            sell_rules = rules['sell_rules']

            fields = buy_rules + sell_rules

        if 'features' in config.keys():
            fields += config['features']

        date_end = datetime.now().strftime('%Y%m%d')
        if 'date_end' in config.keys():
            date_end = config['date_end']

        date_start = None
        if 'date_start' in config.keys():
            date_start = config['date_start']

        df = D.load(config['codes'], start_time=date_start, end_time=date_end, fields=fields, names=fields)

        if 'rules' in config.keys():
            rules = config['rules']
            buy_temp = df[rules['buy_rules']].sum(axis=1)
            df['to_buy'] = buy_temp >= rules['buy_min_count']
            sell_temp = df[rules['sell_rules']].sum(axis=1)
            df['to_sell'] = sell_temp >= rules['sell_min_count']
        df.dropna(inplace=True)
        return df

    def prepare_strategy(self, config:dict):
        algos = config['algos']
        algo_list = [
        ]
        for algo_dict in algos:
            module = importlib.import_module(algo_dict['module'])
            kwargs = None

            algo_class = getattr(module,algo_dict['class'])

            if 'kwargs' in algo_dict.keys():
                kwargs = algo_dict['kwargs']
                algo_obj = algo_class(**kwargs)
            else:
                algo_obj = algo_class()

            algo_list.append(algo_obj)

        s = Strategy(name=config['name'], algo_list=algo_list)
        return s

    def analysis(self, benchmarks=[], task_uids=[]):
        equities = []
        for benchmark in benchmarks:
            bench_df = D.load_code_df(benchmark)
            se = (bench_df['rate']+1).cumprod()
            se.name = benchmark
            equities.append(se)

        for uid in task_uids:
            task = self.tasks[uid]
            filename = get_config_path()+'/results/{}.csv'.format(uid)
            if os.path.exists(filename):
                df = pd.read_csv(filename)
                df['date'] = df['date'].apply(lambda x:str(x))
                df.index = df['date']
                se =(df['rate']+1).cumprod()
                se.name = task['name']
                equities.append(se)

        df_equities = pd.concat(equities,axis=1)
        df_equities.dropna(inplace=True)
        print(df_equities)


        from engine.performance import PerformanceUtils

        df_ratios, df_corr, df_years = PerformanceUtils().calc_equity(df_equity=df_equities)
        print(df_ratios)
        df_equities.plot()
        import matplotlib.pyplot as plt
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
        plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
        plt.show()
        '''
        G.notify(G.GRAPH_RESULT,
                 {
                     'equity': df_equities,
                     'ratio': df_ratios,
                     'corr': df_corr,
                     'yearly': df_years

                 })
        '''
        return df_equities, df_ratios, df_corr, df_years


if __name__ == '__main__':

    task_config = {

    }
    mgr = TaskRunner()
    uid = '708fb7c6-5507-4f81-bb68-55fb9aaeb8ac'
    uid = 'mom_20708fb7c6-5507-4f81-bb68-55fb9aaeb8ac'
    #uid = 'rsrs8fb7c6-5507-4f81-bb68-55fb9aaeb8ac'
    df = mgr.run_by_uid(uid)
    #print(df)
    mgr.analysis(benchmarks=['399006.SZ','000300.SH','000905.SH'], task_uids=[uid])

